#ifndef AST_NODE_H
#define AST_NODE_H

#include <vector>
#include <string>
#include <memory>
#include <iostream>

// 前向声明
class ASTVisitor;

// AST节点类型枚举
enum class NodeType {
    // 程序结构
    PROGRAM,
    COMP_UNIT,
    
    // 声明
    CONST_DECL,
    VAR_DECL,
    CONST_DEF,
    VAR_DEF,
    
    // 函数
    FUNC_DEF,
    FUNC_FPARAM,
    FUNC_FPARAMS,
    FUNC_RPARAM,
    FUNC_RPARAMS,
    
    // 语句
    BLOCK,
    ASSIGN_STMT,
    EXPR_STMT,
    RETURN_STMT,
    IF_STMT,
    
    // 表达式
    BINARY_EXPR,
    UNARY_EXPR,
    CALL_EXPR,
    IDENT_EXPR,
    LITERAL_EXPR,
    PAREN_EXPR,
    
    // 基本类型
    BASIC_TYPE,
    
    // 其他
    INIT_VAL,
    CONST_INIT_VAL
};

// AST节点基类
class ASTNode {
public:
    NodeType nodeType;
    std::vector<std::unique_ptr<ASTNode>> children;
    
    ASTNode(NodeType type) : nodeType(type) {}
    virtual ~ASTNode() = default;
    
    // 访问者模式接口
    virtual void accept(ASTVisitor& visitor) = 0;
    
    // 添加子节点
    void addChild(std::unique_ptr<ASTNode> child) {
        if (child) {
            children.push_back(std::move(child));
        }
    }
    
    // 获取节点类型名称
    std::string getNodeTypeName() const;
    
    // 打印AST树
    void printTree(int depth = 0) const;
};

// 程序节点
class ProgramNode : public ASTNode {
public:
    ProgramNode() : ASTNode(NodeType::PROGRAM) {}
    void accept(ASTVisitor& visitor) override;
};

// 编译单元节点
class CompUnitNode : public ASTNode {
public:
    CompUnitNode() : ASTNode(NodeType::COMP_UNIT) {}
    void accept(ASTVisitor& visitor) override;
};

// 基本类型节点
class BasicTypeNode : public ASTNode {
public:
    std::string typeName; // "int", "float", "void"
    
    BasicTypeNode(const std::string& type) : ASTNode(NodeType::BASIC_TYPE), typeName(type) {}
    void accept(ASTVisitor& visitor) override;
};

// 常量声明节点
class ConstDeclNode : public ASTNode {
public:
    std::unique_ptr<BasicTypeNode> bType;
    
    ConstDeclNode() : ASTNode(NodeType::CONST_DECL) {}
    void accept(ASTVisitor& visitor) override;
};

// 变量声明节点
class VarDeclNode : public ASTNode {
public:
    std::unique_ptr<BasicTypeNode> bType;
    
    VarDeclNode() : ASTNode(NodeType::VAR_DECL) {}
    void accept(ASTVisitor& visitor) override;
};

// 常量定义节点
class ConstDefNode : public ASTNode {
public:
    std::string identifier;
    std::unique_ptr<ASTNode> initVal;
    
    ConstDefNode(const std::string& ident) : ASTNode(NodeType::CONST_DEF), identifier(ident) {}
    void accept(ASTVisitor& visitor) override;
};

// 变量定义节点
class VarDefNode : public ASTNode {
public:
    std::string identifier;
    std::unique_ptr<ASTNode> initVal; // 可选的初始化值
    
    VarDefNode(const std::string& ident) : ASTNode(NodeType::VAR_DEF), identifier(ident) {}
    void accept(ASTVisitor& visitor) override;
};

// 函数定义节点
class FuncDefNode : public ASTNode {
public:
    std::unique_ptr<BasicTypeNode> returnType;
    std::string identifier;
    std::unique_ptr<ASTNode> params; // 参数列表
    std::unique_ptr<ASTNode> body;   // 函数体
    
    FuncDefNode(const std::string& ident) : ASTNode(NodeType::FUNC_DEF), identifier(ident) {}
    void accept(ASTVisitor& visitor) override;
};

// 函数形参节点
class FuncFParamNode : public ASTNode {
public:
    std::unique_ptr<BasicTypeNode> paramType;
    std::string identifier;
    
    FuncFParamNode(const std::string& ident) : ASTNode(NodeType::FUNC_FPARAM), identifier(ident) {}
    void accept(ASTVisitor& visitor) override;
};

// 语句块节点
class BlockNode : public ASTNode {
public:
    BlockNode() : ASTNode(NodeType::BLOCK) {}
    void accept(ASTVisitor& visitor) override;
};

// 赋值语句节点
class AssignStmtNode : public ASTNode {
public:
    std::unique_ptr<ASTNode> lval;
    std::unique_ptr<ASTNode> expr;
    
    AssignStmtNode() : ASTNode(NodeType::ASSIGN_STMT) {}
    void accept(ASTVisitor& visitor) override;
};

// 表达式语句节点
class ExprStmtNode : public ASTNode {
public:
    std::unique_ptr<ASTNode> expr; // 可选的表达式
    
    ExprStmtNode() : ASTNode(NodeType::EXPR_STMT) {}
    void accept(ASTVisitor& visitor) override;
};

// 返回语句节点
class ReturnStmtNode : public ASTNode {
public:
    std::unique_ptr<ASTNode> expr; // 可选的返回表达式
    
    ReturnStmtNode() : ASTNode(NodeType::RETURN_STMT) {}
    void accept(ASTVisitor& visitor) override;
};

// if语句节点
class IfStmtNode : public ASTNode {
public:
    std::unique_ptr<ASTNode> condition;
    std::unique_ptr<ASTNode> thenStmt;
    std::unique_ptr<ASTNode> elseStmt; // 可选的else语句
    
    IfStmtNode() : ASTNode(NodeType::IF_STMT) {}
    void accept(ASTVisitor& visitor) override;
};

// 二元表达式节点
class BinaryExprNode : public ASTNode {
public:
    std::string op; // 运算符
    std::unique_ptr<ASTNode> left;
    std::unique_ptr<ASTNode> right;
    
    BinaryExprNode(const std::string& operator_) : ASTNode(NodeType::BINARY_EXPR), op(operator_) {}
    void accept(ASTVisitor& visitor) override;
};

// 一元表达式节点
class UnaryExprNode : public ASTNode {
public:
    std::string op; // 运算符
    std::unique_ptr<ASTNode> operand;
    
    UnaryExprNode(const std::string& operator_) : ASTNode(NodeType::UNARY_EXPR), op(operator_) {}
    void accept(ASTVisitor& visitor) override;
};

// 函数调用表达式节点
class CallExprNode : public ASTNode {
public:
    std::string funcName;
    std::vector<std::unique_ptr<ASTNode>> args;
    
    CallExprNode(const std::string& name) : ASTNode(NodeType::CALL_EXPR), funcName(name) {}
    void accept(ASTVisitor& visitor) override;
    
    void addArg(std::unique_ptr<ASTNode> arg) {
        if (arg) {
            args.push_back(std::move(arg));
        }
    }
};

// 标识符表达式节点
class IdentExprNode : public ASTNode {
public:
    std::string identifier;
    
    IdentExprNode(const std::string& ident) : ASTNode(NodeType::IDENT_EXPR), identifier(ident) {}
    void accept(ASTVisitor& visitor) override;
};

// 字面量表达式节点
class LiteralExprNode : public ASTNode {
public:
    std::string value;
    std::string literalType; // "int", "float"
    
    LiteralExprNode(const std::string& val, const std::string& type) 
        : ASTNode(NodeType::LITERAL_EXPR), value(val), literalType(type) {}
    void accept(ASTVisitor& visitor) override;
};

// 括号表达式节点
class ParenExprNode : public ASTNode {
public:
    std::unique_ptr<ASTNode> expr;
    
    ParenExprNode() : ASTNode(NodeType::PAREN_EXPR) {}
    void accept(ASTVisitor& visitor) override;
};

// 初始值节点
class InitValNode : public ASTNode {
public:
    std::unique_ptr<ASTNode> expr;
    
    InitValNode() : ASTNode(NodeType::INIT_VAL) {}
    void accept(ASTVisitor& visitor) override;
};

// 常量初始值节点
class ConstInitValNode : public ASTNode {
public:
    std::unique_ptr<ASTNode> expr;
    
    ConstInitValNode() : ASTNode(NodeType::CONST_INIT_VAL) {}
    void accept(ASTVisitor& visitor) override;
};

// 访问者模式基类
class ASTVisitor {
public:
    virtual ~ASTVisitor() = default;
    
    virtual void visit(ProgramNode& node) = 0;
    virtual void visit(CompUnitNode& node) = 0;
    virtual void visit(BasicTypeNode& node) = 0;
    virtual void visit(ConstDeclNode& node) = 0;
    virtual void visit(VarDeclNode& node) = 0;
    virtual void visit(ConstDefNode& node) = 0;
    virtual void visit(VarDefNode& node) = 0;
    virtual void visit(FuncDefNode& node) = 0;
    virtual void visit(FuncFParamNode& node) = 0;
    virtual void visit(BlockNode& node) = 0;
    virtual void visit(AssignStmtNode& node) = 0;
    virtual void visit(ExprStmtNode& node) = 0;
    virtual void visit(ReturnStmtNode& node) = 0;
    virtual void visit(IfStmtNode& node) = 0;
    virtual void visit(BinaryExprNode& node) = 0;
    virtual void visit(UnaryExprNode& node) = 0;
    virtual void visit(CallExprNode& node) = 0;
    virtual void visit(IdentExprNode& node) = 0;
    virtual void visit(LiteralExprNode& node) = 0;
    virtual void visit(ParenExprNode& node) = 0;
    virtual void visit(InitValNode& node) = 0;
    virtual void visit(ConstInitValNode& node) = 0;
};

#endif
